# 对话
import json
from typing import Annotated

from fastapi import APIRouter, HTTPException
from fastapi.params import Path
from fastapi.responses import StreamingResponse
from sqlalchemy import select
from sqlalchemy.orm import Session

from models import User, Chat, Message, engine
from prompts.chat import *
from schemas.chat import (
    UpdateChatRequest,
    GetChatResponse,
    CreateChatRequest,
    UpdateChatResponse,
)
from services.dashscope_tts import qwen_tts_stream
from services.global_analyzer import get_messages_and_analyses
from services.openai_service import stream_generate

router = APIRouter()


@router.post(path="/chat", summary="申请新的对话", tags=["chat"])
async def create_chat(request: CreateChatRequest):
    if request.mode <= 0 or request.mode > 3:
        raise HTTPException(status_code=400, detail="Invalid mode")

    with Session(engine) as session:
        user = session.get(User, request.user_id)
        if user is None:
            raise HTTPException(status_code=404, detail="User not found")

        chat = Chat(user_id=str(user.id), mode=request.mode)
        if request.mode == 2 or request.mode == 3:
            chat.system_prompt = (
                request.situation
            )  # 之后需要换为具体情景对话提示词提示词
        session.add(chat)
        session.commit()
        session.refresh(chat)

    return {"status": 1, "chat_id": chat.id}


@router.put(
    path="/chat/{chat_id}",
    summary="保存新的message",
    tags=["chat"],
    response_model=UpdateChatResponse,
)
async def update_chat(
    chat_id: Annotated[str, Path(description="The ID of conversation, typing UUID4.")],
    request: UpdateChatRequest,
):
    with Session(engine) as session:
        chat = session.get(Chat, chat_id)
        if not chat:
            raise HTTPException(status_code=404, detail="Chat not found")

        messages = chat.messages
        new_message = Message(
            index=len(messages) + 1,
            role=request.role,
            content=request.content,
            chat_id=chat_id,
        )
        session.add(new_message)
        session.commit()
        session.refresh(new_message)

    return UpdateChatResponse(message_id=new_message.id)


@router.get(
    path="/chat/{chat_id}/stream",
    summary="在当前对话中获取回复",
    tags=["chat"],
    response_class=StreamingResponse,
)
async def get_chat(
    chat_id: Annotated[str, Path(title="The ID of conversation, typing UUID4.")],
):
    # history
    messages = []
    with Session(engine) as session:
        # 系统提示词
        chat = session.get(Chat, chat_id)
        if chat.mode == 2 or chat.mode == 3:
            messages.append(
                {
                    "role": "system",
                    "content": prompt_for_situation_chat(chat.system_prompt),
                }
            )
        else:
            messages.append(
                {"role": "system", "content": prompt_for_free_chat()}
            )  # 自由对话模式
        # user & assistant
        statement = (
            select(Message).where(Message.chat_id == chat_id).order_by(Message.index)
        )
        results = session.scalars(statement).all()
        for result in results:
            messages.append({"role": result.role, "content": result.content})

        if messages[-1]["role"] == "assistant":
            raise HTTPException(
                status_code=404,
                detail="最后一次对话以Assistant Message结尾，请检查新对话是否成功上传User Message。",
            )

    async def generate():
        assistant_text = ""
        # 流式处理文本和音频
        sentence_index = 0
        for chunk in stream_generate(messages, sentence_split=True):
            assistant_text += chunk
            yield f"data: {json.dumps({'index': sentence_index,'type': 'text', 'content': chunk})}\n\n"

            for audio_chunk in qwen_tts_stream(chunk):
                yield f"data: {json.dumps({'index': sentence_index, 'type': 'audio', 'content': audio_chunk})}\n\n"

    return StreamingResponse(content=generate(), media_type="text/event-stream")


@router.get(path="/chat/{chat_id}", summary="获取先前的对话记录", tags=["chat"])
async def get_conversation(chat_id: Annotated[str, Path()]):
    try:
        messages = get_messages_and_analyses(chat_id)
    except Exception as e:
        raise HTTPException(
            status_code=500, detail="Error retrieving chat history. Error %s" % str(e)
        )

    return GetChatResponse(chat_id=chat_id, number=len(messages), messages=messages)
